import numpy as np

from core.variable import Variable
from implement.layers.basic.batchnorm import BatchNorm
from implement.layers.basic.model import Model
from implement.layers.basic.affine import Affine
from utils.functions_collect import sigmoid


# 两层网络
# 这种面向对象的方式定义模型的做法（也就是把模型以类为单位进行整合）出自于Chainer框架
class TwoLayersNet(Model):
    def __init__(self, hidden_size, out_size):
        """
        构造函数，初始化两个线性层（全连接层）的实例。

        Parameters:
            out_size (int): 输出层的神经元数量，表示模型输出的维度。
            hidden_size (int): 隐藏层的神经元数量，表示模型中间层的维度。
        """
        super().__init__()
        # 第一个线性层，输入维度为 hidden_size，输出维度为 hidden_size
        self.l1 = Affine(hidden_size)
        # 第二个线性层，输入维度为 hidden_size，输出维度为 out_size
        self.l2 = Affine(out_size)
        self.bn1 = BatchNorm()

    def forward(self, x):
        """
        前向传播函数，描述了数据在神经网络中的正向流动过程。

        Parameters:
            x (Variable): 输入数据，是一个 Variable 对象。

        Returns:
            Variable: 模型的输出，是一个 Variable 对象。
        """
        # 第一层线性层的输出，经过 sigmoid 激活函数
        y = sigmoid(self.bn1(self.l1(x)))

        # 第二层线性层的输出，作为模型的最终输出
        y = self.l2(y)

        return y


if __name__ == '__main__':
    x = Variable(np.random.rand(5, 10), name='x')
    model = TwoLayersNet(100, 10)
    model.plot(x, to_file='two_layers_net.png')
